import random
from pulp import LpProblem, LpMinimize, LpVariable
from datetime import datetime
from argparse import ArgumentParser
from common import format_timedelta


def greedy(universal_set, subsets, costs):
    # U = universal_set
    # S = subsets
    # c(S) = costs
    # F = universal_set
    # C = result
    constructed_set = set()
    result = []
    while len(constructed_set) < len(universal_set):
        optimal_subset_idx = None
        for i in range(len(subsets)):
            if len(subsets[i] - constructed_set) > 0:
                if optimal_subset_idx is None or costs[i] / len(subsets[i] - constructed_set) < costs[optimal_subset_idx] / len(subsets[optimal_subset_idx] - constructed_set):
                    optimal_subset_idx = i
        result.append(optimal_subset_idx)
        constructed_set = constructed_set | subsets[optimal_subset_idx]
    return result


def linear_programming(universal_set, subsets, costs):
    problem = LpProblem('set_covering', LpMinimize)

    variables = []
    for i in range(len(subsets)):
        variables.append(LpVariable('x' + str(i), lowBound=0))

    objective = 0
    for cost, x in zip(costs, variables):
        objective += cost * x
    problem += objective

    f = 0
    for e in universal_set:
        constraint = 0
        frequency = 0
        for i in range(len(subsets)):
            if e in subsets[i]:
                constraint += variables[i]
                frequency += 1
        problem += constraint >= 1
        if frequency > f:
            f = frequency

    status = problem.solve()
    if status != 1:
        raise Exception('status != 1')

    # for i in range(1, 13):
    #     print(problem.constraints['_C' + str(i)])
    # exit()

    result = []
    for i in range(len(subsets)):
        if variables[i].varValue >= 1 / f:
            result.append(i)
    return result


def wrapper(algorithm, universal_set, subsets, costs):
    start_time = datetime.now()
    result = algorithm(universal_set, subsets, costs)
    end_time = datetime.now()
    return result, end_time - start_time


def generate_sample(input_size):
    universal_set = set(range(input_size))

    subsets = []
    if len(universal_set) <= 20:
        subsets.append(universal_set)
    else:
        s0 = set(random.sample(universal_set, 20))
        subsets.append(s0)
        selected = s0
        while len(universal_set - selected) >= 20:
            n = random.randint(1, 20)
            x = random.randint(1, n)
            s = set(random.sample(selected, n - x) + random.sample(universal_set - selected, x))
            subsets.append(s)
            selected = selected | s
        subsets.append(universal_set - selected)

    while len(subsets) < input_size:
        n = random.randint(1, 20)
        n = 20
        subsets.append(set(random.sample(universal_set, n)))

    costs = [1 for _ in range(input_size)]

    return universal_set, subsets, costs


def test_case(algorithm, input_size):
    print('输入规模：%d' % (input_size,))
    universal_set, subsets, costs = generate_sample(input_size)
    result, timedelta = wrapper(algorithm, universal_set, subsets, costs)
    print('用时：' + format_timedelta(timedelta))
    print('子集个数：' + str(len(result)))
    # print(result)
    # for i in result:
    #     print(len(subsets[i]))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--random_seed', type=str, default=19950125)
    parser.add_argument('--algorithm', type=str, default='greedy')
    parser.add_argument('--input_size', type=int, default=100)
    parser.add_argument('--input_file', type=str, default='')
    option = parser.parse_args()

    random.seed(option.random_seed)

    algorithm = {
        'greedy': greedy,
        'lp': linear_programming
    }[option.algorithm]

    if option.input_file.strip() == '':
        test_case(algorithm, option.input_size)
    else:
        subsets = []
        costs = []
        with open(option.input_file, 'r') as f:
            for line in f:
                line = line.strip().split()
                costs.append(float(line[0]))
                subsets.append(set([int(i) for i in line[1:]]))
        universal_set = set()
        for subset in subsets:
            universal_set = universal_set | subset
        result, timedelta = wrapper(algorithm, universal_set, subsets, costs)
        result = sorted([i + 1 for i in result])
        print('输出：' + str(result))
        print('用时：' + format_timedelta(timedelta))
